#' @description  lasso子函数，主要用于计算coef结果
#' @param exprs 表达谱
#' @param clinical 临床信息
#' @param genes 目标基因
#' @param seed 设置种子
#' @export
compute_lasso_coef <- function(
    exprs = NULL, clinical = NULL, genes = NULL,
    seed = 1110, time_col = "time", status_col = "status", ...) {
    suppressPackageStartupMessages(library(cowplot))
    suppressPackageStartupMessages(library(tidyverse))
    suppressPackageStartupMessages(library(survival))
    suppressPackageStartupMessages(library(ggpubr))

    exprs_part <- exprs[genes, ]
    clinical <- clinical[na.omit(match(colnames(exprs_part), clinical$sample)), ]
    exprs_part <- exprs_part[, na.omit(match(clinical$sample, colnames(exprs_part)))]

    colnames(clinical)[which(colnames(clinical) == time_col)] <- "time"
    colnames(clinical)[which(colnames(clinical) == status_col)] <- "status"

    clinical %<>% filter(time > 0)

    lasso_res <- vector("list", 3)
    suppressPackageStartupMessages(library(glmnet))
    y <- data.matrix(Surv(as.numeric(clinical$time), as.numeric(clinical$status)))
    x <- exprs_part %>%
        t() %>%
        as.data.frame() %>%
        as.matrix()

    set.seed(seed)
    fit <- glmnet(x, y, family = "cox", alpha = 1)
    lasso_res[[2]] <- fit

    cvfit <- cv.glmnet(x, y, family = "cox")
    lasso_res[[3]] <- cvfit

    tmp <- coef(object = cvfit, s = "lambda.min")

    Signature_Coef_min_os <- tmp %>%
        as.matrix() %>%
        as.data.frame() %>%
        rownames_to_column("gene") %>%
        dplyr::rename(coef = 2) %>%
        filter(coef != 0)

    colnames(Signature_Coef_min_os) <- c("symbol", "coef")

    lasso_res[[1]] <- Signature_Coef_min_os

    return(lasso_res)
}

#' @description 运行lasso方法进行构建模型
#' @export
lasso_plot <- \(lasso_res){
    ss <<- lasso_res
    p1 <- ~ {
        plot(ss[[2]], xvar = c("lambda"))
        abline(v = log(ss[[3]][["lambda.min"]]), lty = "dotted")
    }

    p2 <- ~ {
        plot(ss[[3]])
    }

    coef_bar_p <- ggplot(
        data = ss[[1]] %>%
            mutate(
                group = ifelse(.$coef > 0, "1", "2"),
                min_coef = min(coef), max_coef = max(coef)
            ),
        mapping = aes(x = coef, y = factor(reorder(symbol, coef)), fill = factor(group))
    ) +
        geom_bar(stat = "identity", color = "black", width = .8) +
        ggpubr::theme_pubclean(13) +
        guides(fill = "none") +
        ylab("Gene") +
        xlab("Coefficients") +
        scale_fill_manual(values = c("2" = "#3171a5", "1" = "#FF7F00"))

    p <- plot_grid(p1, p2, coef_bar_p, labels = "AUTO", rel_widths = c(1.2, 1.2, 1), nrow = 1)
    p
}

#' @description 多因素cox方法进行构建模型
#' @export
multi_cox_model <- function(
    exprs = NULL, clinical = NULL, genes = NULL,
    time_col = "time", status_col = "status", ...) {
    require(survival)
    suppressPackageStartupMessages(require(survminer))
	
    exprs_part <- exprs[genes, ]
    exprs_part <- as.data.frame(t(exprs_part))
    clinical <- clinical[na.omit(match(rownames(exprs_part), clinical[, "sample"])), ]
    exprs_part <- exprs_part[na.omit(match(clinical[, "sample"], rownames(exprs_part))), ]

    colnames(clinical)[which(colnames(clinical) == time_col)] <- "time"
    colnames(clinical)[which(colnames(clinical) == status_col)] <- "status"

    multicox_formulas <-
        as.formula(paste(
            str_glue("Surv(clinical$time, clinical$status)~"),
            paste0(sep = "`", colnames(exprs_part)[1:ncol(exprs_part)], sep = "`", collapse = "+")
        ))
    result_0 <- coxph(
        formula = multicox_formulas,
        data = exprs_part
    )

    result <- summary(result_0)

    result2 <- data.frame(
        symbol = rownames(result$coefficients),
        coef = result$coefficients[, 1],
        Hazard_Ratio = result$coefficients[, 2],
        p.value = result$coefficients[, 5],
        lower_.95 = result$conf.int[, "lower .95"],
        upper_.95 = result$conf.int[, "upper .95"]
    )

    result2 <- result2 %>%
        mutate(HR = paste0(
            round(Hazard_Ratio, 2),
            "(",
            round(lower_.95, 2),
            "-",
            round(upper_.95, 2),
            ")"
        ))

    temp <- list()
    temp[[1]] <- result2
    temp[[2]] <- result_0

    p <- survminer::ggforest(
        model = temp[[2]],
        data = clinical,
        fontsize = 1,
        main = "Multivariable Analysis"
    )

    temp[[3]] <- p
    names(temp) <- c("MUltiCox_res_df", "MUltiCox_res", "MultiCox_plot")

    return(temp)
}

#' @description  lasso子函数，主要用于计算coef结果
#' @param exprs 表达谱
#' @param clinical 临床信息
#' @param genes 目标基因
#' @param seed 设置种子
#' @export
model_lst <- function(model_type, ...) {
    if (!{
        model_type %in% c("lasso", "multi_cox", "PCA1_minus_PCA2", "PCA1_sum_PCA2")
    }) {
        fun_defined <- \(...) {
            lst_inner <- list(...)
            for (i in names(lst_inner)) {
                if (i != "") {
                    assign(i, lst_inner[[i]])
                }
            }
            fun_defined <- \(...){}

            body(fun_defined) <- body(lst_inner[[which(map_lgl(lst_inner, ~ is.function(.x)))]])
            return(list("defined" = fun_defined()))
        }
        return(fun_defined)
    }

    model_inner <- list(
        lasso = \(x, ...) {
            if (!any(grepl(pattern = "gene", x = colnames(x)))) {
                cli::cli_alert_danger("单因素cox结果中不存在{.var gene}列, lasso报错")
            }

            lst_inner <- list(...)
            for (i in names(lst_inner)) {
                if (i != "") {
                    assign(i, lst_inner[[i]])
                }
            }

            genes <- x$gene

            lasso_res <- compute_lasso_coef(
                exprs = train_data$tumor_exprs,
                clinical = train_clinical,
                genes = genes, seed = 1110,
                time_col = time_col,
                status_col = status_col
            )

            p <- suppressMessages(lasso_plot(lasso_res))
            res <- rlist::list.append(lasso_res, p)
        },
        multi_cox = \(x, ...) {
            if (!any(grepl(pattern = "gene", x = colnames(x)))) {
                cli::cli_alert_danger("单因素cox结果中不存在{.var gene}列")
            }

            lst_inner <- list(...)
            for (i in names(lst_inner)) {
                if (i != "") {
                    assign(i, lst_inner[[i]])
                }
            }

            genes <- x$gene

            res <- multi_cox_model(
                exprs = train_data$tumor_exprs,
                clinical = train_clinical,
                genes = genes, seed = 1110,
                time_col = time_col,
                status_col = status_col
            )
            res
        },
        PCA1_minus_PCA2 = \(x, ...) {
            if (!any(grepl(pattern = "gene", x = colnames(x)))) {
                cli::cli_alert_danger("单因素cox结果中不存在{.var gene}列")
            }
            if (!any(grepl(pattern = "HR", x = colnames(x)))) {
                cli::cli_alert_danger("单因素cox结果中不存在{.var HR}列")
            }

            up_gene <- x %>%
                filter(HR > 1) %>%
                pull(gene)
            down_gene <- x %>%
                filter(HR < 1) %>%
                pull(gene)

            x_lst <- vector("list", 2)
            pacman::p_load(FactoMineR)

            if (!exists("train_data")) {
                cli::cli_alert_danger("在方法{.var PCA1_minus_PCA2}中，变量{.var train_data}不存在，请加载。")
                return()
            }
            exprs <- train_data$tumor_exprs

            x_lst[[1]] <- intersect(x_lst[[1]], rownames(train_data$tumor_exprs))
            x_lst[[2]] <- intersect(x_lst[[2]], rownames(exprs))

            if (length(x_lst[[1]]) <= 2 || length(x_lst[[2]] <= 2)) {
                cli::cli_alert_danger("在方法{.var PCA1_minus_PCA2}中，输入的基因列表的第一个或第二个元素在表达谱中交集小于2个，程序跳出。")
                return()
            }

            pcaH <- PCA(X = exprs[x_lst[[1]], ] %>% t(), graph = FALSE, ncp = 2)
            pcaL <- PCA(X = exprs[x_lst[[2]], ] %>% t(), graph = FALSE, ncp = 2)
            score <- pcaH$ind$contrib[, 1] + pcaH$ind$contrib[, 2] - pcaL$ind$contrib[, 1] - pcaL$ind$contrib[, 2]
            pca_res <- data.frame(PC1PC2 = score) %>% t()
            list("pca_res" = pca_res)
        },
        PCA1_sum_PCA2 = \(x, ...) {
            if (!any(grepl(pattern = "gene", x = colnames(x)))) {
                cli::cli_alert_danger("单因素cox结果中不存在{.var gene}列")
            }

            if (!exists("train_data")) {
                cli::cli_alert_danger("在方法{.var PCA1_minus_PCA2}中，变量{.var train_data}不存在，请加载。")
            }
            exprs <- train_data$tumor_exprs
            x <- x$gene

            x <- intersect(x, rownames(exprs))
            pca_res <- prcomp(x = exprs[x, ], retx = F, scale = T, center = F)
            pca_res <- pca_res$rotation %>%
                as.data.frame() %>%
                transmute(PC1PC2 = PC1 + PC2) %>%
                rename(riskscore = 1) %>%
                t() %>%
                as.data.frame()

            return(list("pca_res" = pca_res))
        }
    )

    return(model_inner[[model_type]])
}

# model_lst <- sapply(model_lst, \(x) safely(x))

#' @description 用于计算score
#' @param exprs 表达谱
#' @param model_coef 模型使用coef信息
#' @details 内部函数，不用于外部调用
compute_score_in_wgcan_p <- function(exprs, model_coef = NULL) {
    exprs <- exprs[model_coef[, 1], , drop = F] %>%
        t() %>%
        as.data.frame()

    miss_gene <- model_coef[, 1][which(is.na(exprs[1, ]))]
    miss_gene_symbol <- paste0(miss_gene, collapse = ", ")

    exprs[is.na(exprs)] <- 0

    riskscore_res <- as.matrix(exprs) %*% as.matrix(model_coef[, 2, drop = F]) %>%
        as.data.frame() %>%
        dplyr::rename(riskscore = 1)

    if (length(miss_gene) == 1) {
        cli::cli_alert_info(str_glue("'WARNING:' {miss_gene_symbol} 在当前队列的表达谱中没有匹配到。\n"))
    }

    return(riskscore_res)
}

#' @description  lasso子函数，主要用于计算coef结果
#' @param model 字符串，建模方式
#' @param x 不同建模方式所使用的基因系数或者基因list等等
#' @param exprs 表达谱
#' @details 参考{/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/NewLover/immunetherapy_FUN.R}中的`get_score`说明
#' @export
get_score_in_wgcan_p <- function(model, x, exprs, ...) {
    if (!{
        model %in% c("lasso", "multi_cox", "PCA1_minus_PCA2", "PCA1_sum_PCA2")
    }) {
        score_res <- compute_score_in_wgcan_p(exprs = x, model_coef = data.frame(gene = "riskscore", coef = 1))
        return(score_res)
    }

    coef_1_df <- data.frame(gene = "riskscore", coef = 1)

    model_fun <- list(
        lasso = \(...){
            if (!is.data.frame(x)) {
                cli::cli_alert_danger(str_glue("{model}要求输入x是DF，且第一列为基因名，第二列为相关系数"))
                return()
            }
            score_res <- compute_score_in_wgcan_p(exprs, model_coef = x)
            return(score_res)
        },
        multi_cox = \(...) {
            if (!is.data.frame(x)) {
                cli::cli_alert_danger(str_glue("{model}要求输入x是DF，且第一列为基因名，第二列为相关系数"))
                return()
            }
            score_res <- compute_score_in_wgcan_p(exprs, model_coef = x)
            return(score_res)
        },
        PCA1_minus_PCA2 = \(...) {
            score_res <- compute_score_in_wgcan_p(exprs = x, coef_1_df)
            return(score_res)
        },
        PCA1_sum_PCA2 = \(...) {
            score_res <- compute_score_in_wgcan_p(exprs = x, coef_1_df)
            return(score_res)
        }
    )

    return(model_fun[[model]]())
}
